# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
"""
Implements the Generalized R-CNN framework
"""
import torch
from torch import nn
from maskrcnn_benchmark.structures.image_list import to_image_list, ImageList
from maskrcnn_benchmark.data.collate_batch import ImageList5D
from ..backbone import build_backbone
from ..rpn.rpn import build_rpn
from ..roi_heads.roi_heads import build_roi_heads


class GeneralizedRCNN(nn.Module):
    """
    Main class for Generalized R-CNN. Currently supports boxes and masks.
    It consists of three main parts:
    - backbone
    - rpn
    - heads: takes the features + the proposals from the RPN and computes
        detections / masks from it.
    """

    def __init__(self, cfg):
        super(GeneralizedRCNN, self).__init__()
        self.cfg = cfg
        self.backbone = build_backbone(cfg)
        self.rpn = build_rpn(cfg, self.backbone.out_channels)
        self.roi_heads = build_roi_heads(cfg, self.backbone.out_channels)

    def forward(self, images, targets=None, **kwargs):
        """
        Arguments:
            images (list[Tensor] or ImageList): images to be processed
            targets (list[BoxList]): ground-truth boxes present in the image (optional)

        Returns:
            result (list[BoxList] or dict[Tensor]): the output from the model.
                During training, it returns a dict[Tensor] which contains the losses.
                During testing, it returns list[BoxList] contains additional fields
                like `scores`, `labels` and `mask` (for Mask R-CNN models).

        """
        if self.training and targets is None:
            raise ValueError("In training mode, targets should be passed")
        
        if isinstance(images, ImageList5D):
            # ---- Normalize to 5D (B,A,C,H,W) even if some batch came as 4D ----
            t = images.tensors
            if t.dim() == 4:
                # Prefer recovering B and A from image_sizes_ba; fallback to A=1
                if getattr(images, "image_sizes_ba", None):
                    B = len(images.image_sizes_ba)
                    A = len(images.image_sizes_ba[0]) if B > 0 else 1
                    BA, C, H, W = t.shape
                    if B * A == BA:
                        t = t.view(B, A, C, H, W)
                    else:
                        # fallback: treat as A=1
                        B = BA
                        t = t.view(B, 1, C, H, W)
                else:
                    # no sizes available → A=1
                    B, C, H, W = t.shape
                    t = t.view(B, 1, C, H, W)
                images.tensors = t
            B, A, C, H, W = images.tensors.shape

            # 1) 构造每个视图的标准 ImageList（不展平 BA）
            views = []
            for a in range(A):
                tensors_4d = images.tensors[:, a, :, :, :]
                if getattr(images, "image_sizes_ba", None):
                    sizes_a = [tuple(images.image_sizes_ba[b][a]) for b in range(B)]
                else:
                    sizes_a = [(H, W)] * B
                views.append(ImageList(tensors_4d, sizes_a))

            # 2) anchor 视图（默认 0）走 RPN / ROIHeads baseline
            anchor = self.cfg.MODEL.DIVALIGN.ANCHOR_VIEW if hasattr(self, "cfg") else 0
            images_anchor = to_image_list(views[anchor])
            # ---- 只取 anchor 视图的 targets（不做空 GT 过滤）----
            targets_anchor = None
            if targets is not None:
                if isinstance(targets[0], (list, tuple)):
                    # 形如 [[BoxList_v0, BoxList_v1], [ ... ], ...] -> 取 anchor 那份
                    targets_anchor = [t_group[anchor] for t_group in targets]
                else:
                    # 已是一维列表（可能是已展平成 BA），取前 B 个
                    B = images_anchor.tensors.size(0)
                    targets_anchor = targets[:B]
            feats_anchor  = self.backbone(images_anchor.tensors)
            proposals, proposal_losses = self.rpn(images_anchor, feats_anchor, targets_anchor)   
            # === 新增：为每个视图准备对应尺寸的 proposals（把锚视图 proposals 缩放到该视图的尺寸） ===
            proposals_views = None
            if self.training and A > 1:
                proposals_views = []
                for a, v in enumerate(views):
                    if a == anchor:
                        # 锚视图直接复用 RPN 的 proposals
                        proposals_views.append(proposals)
                    else:
                        # 将每张图的 proposals 按 (W_a,H_a) -> (W_v,H_v) 线性缩放
                        mapped_list = []
                        for b, prop in enumerate(proposals):
                            # ImageList 里保存的是 (H,W)，而 BoxList.resize 期望 (W,H)
                            H_v, W_v = v.image_sizes[b]
                            mapped_list.append(prop.resize((W_v, H_v)))
                        proposals_views.append(mapped_list)
             # 3) baseline roi heads
            x, result, detector_losses = self.roi_heads(feats_anchor, proposals, targets_anchor)

            # 4) 训练阶段：其其余视图仅做 backbone，供对齐损失（使用各自视图的 proposals 做池化，避免越界）
            if self.training and A > 1:
                features_views = [feats_anchor]
                for a, v in enumerate(views):
                    if a == anchor:
                        continue
                    feats_a = self.backbone(v.tensors)
                    features_views.append(feats_a)
                # 只取对齐损失（CombinedROIHeads 会合并到返回的 losses 里）
                _x2, _det2, extra_align_losses = self.roi_heads(feats_anchor, proposals, targets_anchor, features_views=features_views, proposals_views=proposals_views)
                detector_losses.update({k: v for k, v in extra_align_losses.items() if "align" in k})

            if self.training:
                losses = {}
                losses.update(detector_losses)
                losses.update(proposal_losses)
                return losses
            return result
        
        # ---- 非 ImageList5D 的老逻辑保持不变 ----
        images = to_image_list(images)
        features = self.backbone(images.tensors)
        proposals, proposal_losses = self.rpn(images, features, targets)
        if self.roi_heads:
            x, result, detector_losses = self.roi_heads(features, proposals, targets)
        else:
            # RPN-only models don't have roi_heads
            x = features
            result = proposals
            detector_losses = {}

        if self.training:
            losses = {}
            losses.update(detector_losses)
            losses.update(proposal_losses)
            return losses
        return result
